package com.bz.service.npl;

import opennlp.tools.namefind.TokenNameFinderModel;
import opennlp.tools.namefind.NameFinderME;
import opennlp.tools.tokenize.TokenizerModel;
import opennlp.tools.tokenize.TokenizerME;
import opennlp.tools.util.Span;

import java.io.InputStream;
import java.util.Arrays;

public class NlpService {

    private final TokenizerME tokenizer;
    private final NameFinderME nameFinder;

    public NlpService() throws Exception {
        // 加载分词模型
        try (InputStream tokenModelIn = getClass().getResourceAsStream("/models/en-token.bin")) {
            TokenizerModel tokenModel = new TokenizerModel(tokenModelIn);
            this.tokenizer = new TokenizerME(tokenModel);
        }

        // 加载命名实体识别模型
        try (InputStream nameModelIn = getClass().getResourceAsStream("/models/en-ner-person.bin")) {
            TokenNameFinderModel nameModel = new TokenNameFinderModel(nameModelIn);
            this.nameFinder = new NameFinderME(nameModel);
        }
    }

    public String[] tokenize(String sentence) {
        return tokenizer.tokenize(sentence);
    }

    public String[] findNames(String sentence) {
        String[] tokens = tokenize(sentence);
        Span[] nameSpans = nameFinder.find(tokens);
        return Arrays.stream(nameSpans)
                .map(span -> String.join(" ", Arrays.copyOfRange(tokens, span.getStart(), span.getEnd())))
                .toArray(String[]::new);
    }
}